import sys 
sys.path.append('/home/aistudio/external-libraries')

import numpy as np
import meshio
from scipy.spatial import Delaunay
from meshpy.tet import MeshInfo, build, Options
from restore_model_tools import restore_model, plot_truck_flow_field, standardize_NS_Equations_3D, data
from mesh_refinement_tools import plot_truck_add_points_3d, get_top_k_indices, tetrahedron_centroid, tetrahedron_midpoint, remove_duplicate_points, remove_duplicate_rows, find_left_points, find_points_on_truck, convert_small_numbers_to_zero

# 对哪个level的网格进行加密？
level = 1

# 加载模型
truck_model = restore_model(model_path = "../2_PINN_training/Model/model-75000.pdparams")
plot_truck_flow_field(truck_model)

# Get the node coordinates
mesh = meshio.read('../0_Mesh_files/level'+str(level)+'_truck.vtk', file_format='vtk')
vertices = mesh.points
tetrahedrons = mesh.cells[0].data

print("\n待加密的level"+str(level)+"网格信息:\n")
print("-原网格节点个数:", vertices.shape[0])
print(vertices)
print("\n-原网格四面体个数:", tetrahedrons.shape[0])
print(tetrahedrons)

# 计算每个四面体的质心
centroid_points = tetrahedron_centroid(tetrahedrons, vertices)

# 找到残差较大的质心点
X = centroid_points
standardize_X = data.transform_inputs(X)
[f1, f2, f3, f4] = truck_model.predict(standardize_X, operator=standardize_NS_Equations_3D)
err_eq = np.absolute(f1) + np.absolute(f2) + np.absolute(f3) + np.absolute(f4)
x_id = get_top_k_indices(err_eq, k=500)
add_points = X[x_id]
plot_truck_add_points_3d(add_points)

# 计算需要被细化的单元的Edge中点
tets_id = x_id
refine_tets = tetrahedrons[tets_id]
refine_midpoints = tetrahedron_midpoint(refine_tets, vertices)
refine_midpoints = remove_duplicate_points(refine_midpoints)
print("待加密四面体网格个数:", refine_tets.shape)
print("待加密四面体边中点个数:", refine_midpoints.shape)
plot_truck_add_points_3d(refine_midpoints)

# 新网格的所有节点
new_mesh_points = convert_small_numbers_to_zero(np.vstack((vertices, refine_midpoints)), threshold = 1e-8)
print("新网格的所有网格节点数：", new_mesh_points.shape[0])
# 内边界点、外边界点和内部点(内外边界点有重合的点)
inner_points = find_points_on_truck(new_mesh_points)
print("新网格汽车表面的网格节点数：", inner_points.shape[0])
# 删除外边界中重合的点
other_points = find_left_points(new_mesh_points, inner_points)
# 拼接
all_mesh_points = np.vstack((inner_points, other_points))

# 四面体化，生成内外两个凸包，然后合并并删除重复的面
outer_tet = Delaunay(all_mesh_points, furthest_site=False, incremental=False, qhull_options=None)
inner_tet = Delaunay(inner_points, furthest_site=False, incremental=False, qhull_options=None)

wrong_convex_hull = np.vstack((inner_tet.convex_hull, outer_tet.convex_hull))
right_convex_hull = remove_duplicate_rows(wrong_convex_hull)

# 对加密后的原有网格节点进行四面体化并输出VTK文件
refined_mesh_info = MeshInfo()
refined_mesh_info.set_points(outer_tet.points)
refined_mesh_info.set_facets(right_convex_hull)
refined_mesh_info.set_holes([(2, 85, 19.25)])
refined_mesh = build(refined_mesh_info, options=Options("p"))
new_mesh = "level" + str(level + 1) + "_refined_truck.vtk"
refined_mesh.write_vtk(new_mesh)




